package com.ai.controller;

import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.time.LocalDateTime;

@RestController
public class StreamController {

    /**
     * 流式输出接口
     */
    @GetMapping(value = "/stream/print", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamPrint() {
        return Flux.interval(Duration.ofSeconds(1))
                .take(10) // 限制输出10条
                .map(sequence -> {
                    String message = "Stream message " + sequence + " at " + LocalDateTime.now();
                    return "springboot" + message + "\n\n";
                });
    }

    /**
     * 模拟长时间任务流式输出
     */
    @GetMapping(value = "/stream/long-task", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> longTask() {
        return Flux.create(fluxSink -> {
            try {
                for (int i = 1; i <= 10; i++) {
                    String step = "Processing step " + i + "/10...";
                    fluxSink.next("springboot" + step + "\n\n");

                    // 模拟处理时间
                    Thread.sleep(1000);

                    if (i == 10) {
                        fluxSink.next("springboot Task completed successfully!\n\n");
                    }
                }
                fluxSink.complete();
            } catch (InterruptedException e) {
                fluxSink.error(e);
            }
        });
    }

    /**
     * 自定义消息的流式输出
     */
    @PostMapping(value = "/stream/custom", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> customStream(@RequestBody StreamRequest request) {
        return Flux.range(1, request.getCount())
                .delayElements(Duration.ofMillis(request.getDelay()))
                .map(i -> {
                    String message = String.format(request.getMessage(), i);
                    return "stream " + message + "\n\n";
                });
    }

    public static class StreamRequest {
        private String message;
        private int count;
        private long delay;

        // getter和setter
        public String getMessage() { return message; }
        public void setMessage(String message) { this.message = message; }
        public int getCount() { return count; }
        public void setCount(int count) { this.count = count; }
        public long getDelay() { return delay; }
        public void setDelay(long delay) { this.delay = delay; }
    }
}